'''
This script is used to transfer the json format in annotations of SODFormer to the npy format, which is often used in recent works.
the struction of the outcomes
shape->(N,) N denotes the number of the bbox
each bbox-> (t, x, y, w, h, class_id) -> (uint64, float32, float32, float32, float32, uint8)
Author: Hatins
Time: 2024/03/12
The output will be added at the h5_input_path.
We move the labels which have not the corresponding aps images
'''

import numpy as np
import os
import json
import h5py
import hdf5plugin
import argparse 
from tqdm import tqdm

# import ipdb

pku_scenes = ['normal']
pku_splits = ['train', 'val']

event_widht = 1280
event_height = 720

def get_number_in_str(file_name):
    numeric_part = ''.join(char for char in file_name if char.isnumeric())
    return int(numeric_part)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--txt_input_path', type = str, default='/data2/zht/nuclear_v6/annotations/',help='josn label path')
    parser.add_argument('--h5_input_path', type = str, default='/data2/zht/nuclear_v6/nuclear_H5',help='raw event path')
    args = parser.parse_args()

    delete_num = 0
    for mode in pku_splits:
        each_mode_josn_dict_path = args.txt_input_path + '/' + mode
        each_mode_h5_dict_path = args.h5_input_path + '/' + mode
        for scene in pku_scenes:
            each_mode_scene_josn_dict_path = each_mode_josn_dict_path + '/' + scene
            each_mode_scene_h5_dict_path = each_mode_h5_dict_path + '/' + scene
            print('prcessing the '+ scene + ' scene in ' + mode)
            for filename in tqdm(os.listdir(each_mode_scene_josn_dict_path)):
    
                josn_file_path = os.path.join(each_mode_scene_josn_dict_path, filename)
                josn_file_path_ = josn_file_path + '/annotations'
                if os.path.exists(josn_file_path_):
                    josn_file_path = josn_file_path_
                else:
                    josn_file_path = josn_file_path + '/annotation'
                h5_file_path = each_mode_scene_h5_dict_path + '/' + filename + '.raw.h5'
                init_time = each_mode_scene_h5_dict_path + '/' + filename + '.raw.init_time.h5'
                save_path = each_mode_scene_h5_dict_path + '/' + filename
                file_list = os.listdir(josn_file_path)
                # npy_file = np.load(npy_path)
                if not os.path.exists(h5_file_path):
                    continue
                print(os.path.exists(h5_file_path))
                h5_file = h5py.File(h5_file_path)
                init_time = h5py.File(init_time)

                # frames_info = h5_file['frames']
                # timestamp_info = frames_info['timestamp']

                init_time = init_time['init_time']['time'][()]
     
                annotations_list = []
                for josn_file_name in file_list:

                    if josn_file_name == 'classes.txt':
                        continue
                    ann_data_list = []
                    file_path = os.path.join(josn_file_path, josn_file_name)
              
                    with open(file_path, 'r') as file:
                        lines = file.readlines()
                        for line in lines:
                            print(file_path)
                            ann_data = [float(num_str) for num_str in line.strip().split()]
            
                            ann_data_list.append(ann_data)

                        ann_data_list = np.array(ann_data_list)
              
                        bbox_num = len(ann_data_list)

                        timestamp = josn_file_name.split(".")[0]
                        timestamp = int(timestamp.split("_")[-1])

                        timestamp =  timestamp - init_time

                        for bbox_index in range(bbox_num):
                            x_tl = (ann_data_list[bbox_index][1] - 0.5 * ann_data_list[bbox_index][3]) * event_widht
                            y_tl = (ann_data_list[bbox_index][2] - 0.5 * ann_data_list[bbox_index][4]) * event_height

                            width = ann_data_list[bbox_index][3] * event_widht
                            height = ann_data_list[bbox_index][4] * event_height

                            if width == 0 or height == 0:
                                print(file_path, ' has the width or height = 0')

                            class_id = ann_data_list[bbox_index][0]
                            annotation = (timestamp, x_tl, y_tl, width, height, class_id)
                            annotations_list.append(annotation)


                dtype = [('t', np.uint64), ('x', np.float32), ('y', np.float32), ('w', np.float32), ('h', np.float32),  ('class_id', np.uint8)]

                annotations_list_numpy = np.array(annotations_list, dtype = dtype)
                sorted_data_array = np.sort(annotations_list_numpy, order='t')
                if len(annotations_list_numpy) == 0:
                    deleted_file = save_path + '.raw.h5'
                    if os.path.exists(deleted_file):
                        os.remove(deleted_file)
                else:
                    save_path = save_path + '_bbox'
                    np.save(save_path, sorted_data_array)
    print('delete {} wrong labels'.format(delete_num))


